#include <rclcpp/rclcpp.hpp>
#include <geometry_msgs/msg/twist_stamped.hpp>
#include <geometry_msgs/msg/pose_stamped.hpp>
#include <tf2_ros/transform_listener.h>
#include <tf2_ros/buffer.h>
#include <eigen3/Eigen/Dense>

class ForceControlNode : public rclcpp::Node
{
public:
    ForceControlNode() : Node("force_control_node")
    {
        // 初始化参数
        M_position_ = Eigen::Matrix3d::Identity() * 280.0;      // 质量矩阵，示例为单位矩阵
        K_position_ = Eigen::Matrix3d::Identity() * 1.0;       // 刚度矩阵，增益设为1
        B_position_ = Eigen::Matrix3d::Identity() * 50.0;     // 阻尼矩阵，增益设为10
        force_torque_reference_ = Eigen::Vector3d(0.0, 0.0, 0.0); // 参考力，设为零
        position_reference_ = Eigen::Vector3d(0.0, 0.0, 0.0);     // 参考位置，设为零
        velocity_ = Eigen::Vector3d::Zero();                      // 初始速度为零
        dt_ = 0.01;                                               // 时间步长，100Hz
        current_force_ = Eigen::Vector3d::Zero();                 // 初始化当前力

        // 创建订阅者，订阅力传感器数据
        force_sub_ = create_subscription<geometry_msgs::msg::TwistStamped>(
            "/force_data", 10, std::bind(&ForceControlNode::force_callback, this, std::placeholders::_1));

        // 初始化TF监听器
        tf_buffer_ = std::make_shared<tf2_ros::Buffer>(this->get_clock());
        tf_listener_ = std::make_shared<tf2_ros::TransformListener>(*tf_buffer_);

        // 创建发布者，发布位置指令到moveit_servo
        pose_pub_ = create_publisher<geometry_msgs::msg::PoseStamped>("/servo_node/delta_pose_cmds", 10);

        // 创建定时器，100Hz运行控制循环
        timer_ = create_wall_timer(std::chrono::milliseconds(10), std::bind(&ForceControlNode::control_loop, this));

        RCLCPP_INFO(this->get_logger(), "Force control node initialized.");
    }

private:
    void force_callback(const geometry_msgs::msg::TwistStamped::SharedPtr msg)
    {
        // 从TwistStamped消息中提取当前力
        current_force_.x() = msg->twist.linear.x;
        current_force_.y() = msg->twist.linear.y;
        current_force_.z() = msg->twist.linear.z;
    }

    void control_loop()
    {
        // 获取当前TCP位置（通过TF变换）
        geometry_msgs::msg::TransformStamped transform;
        try {
            transform = tf_buffer_->lookupTransform("base_link", "end", tf2::TimePointZero);
            current_position_.x() = transform.transform.translation.x;
            current_position_.y() = transform.transform.translation.y;
            current_position_.z() = transform.transform.translation.z;
        } catch (tf2::TransformException &ex) {
            RCLCPP_WARN(this->get_logger(), "无法获取TF变换: %s", ex.what());
            return;
        }

        // 计算力差
        Eigen::Vector3d force_difference = current_force_ - force_torque_reference_;

        // 计算位置差
        Eigen::Vector3d position_difference = current_position_ - position_reference_;

        // 计算刚度项和阻尼项
        Eigen::Vector3d kp_position_term = K_position_ * position_difference;
        Eigen::Vector3d b_velocity_term = B_position_ * velocity_;

        // 计算加速度
        Eigen::Vector3d acceleration = M_position_.inverse() * (force_difference - kp_position_term - b_velocity_term);

        // 更新速度
        velocity_ += acceleration * dt_;

        // 计算位置增量
        Eigen::Vector3d delta_position = velocity_ * dt_;

        // 计算新的位置
        Eigen::Vector3d new_position = current_position_ - delta_position;

        // 创建并填充PoseStamped消息
        geometry_msgs::msg::PoseStamped pose_cmd;
        pose_cmd.header.stamp = this->get_clock()->now();
        pose_cmd.header.frame_id = "base_link";
        pose_cmd.pose.position.x = new_position.x();
        pose_cmd.pose.position.y = new_position.y();
        pose_cmd.pose.position.z = new_position.z();
        // 保持当前朝向不变
        pose_cmd.pose.orientation = transform.transform.rotation;

        // 发布位置指令
        pose_pub_->publish(pose_cmd);

        // 日志输出，便于调试
        RCLCPP_INFO(this->get_logger(), "New position: x=%.3f, y=%.3f, z=%.3f", 
                    new_position.x(), new_position.y(), new_position.z());
    }

    // 成员变量
    rclcpp::Subscription<geometry_msgs::msg::TwistStamped>::SharedPtr force_sub_;
    std::shared_ptr<tf2_ros::Buffer> tf_buffer_;
    std::shared_ptr<tf2_ros::TransformListener> tf_listener_;
    rclcpp::Publisher<geometry_msgs::msg::PoseStamped>::SharedPtr pose_pub_;
    rclcpp::TimerBase::SharedPtr timer_;

    Eigen::Vector3d current_force_;          // 当前力
    Eigen::Vector3d current_position_;       // 当前位置
    Eigen::Vector3d velocity_;               // 速度
    Eigen::Matrix3d M_position_;             // 质量矩阵
    Eigen::Matrix3d K_position_;             // 刚度矩阵
    Eigen::Matrix3d B_position_;             // 阻尼矩阵
    Eigen::Vector3d force_torque_reference_; // 参考力
    Eigen::Vector3d position_reference_;     // 参考位置
    double dt_;                              // 时间步长
};

int main(int argc, char **argv)
{
    rclcpp::init(argc, argv);
    auto node = std::make_shared<ForceControlNode>();
    rclcpp::spin(node);
    rclcpp::shutdown();
    return 0;
}